#include "pkcs11/pkcs11test.h"
#include "cryptoki_ext.h"

// #include <nss/pkcs11n.h>

typedef CK_RV (*CK_NSS_GetFIPSStatus)(CK_SESSION_HANDLE hSession,
                                      CK_OBJECT_HANDLE hObject,
                                      CK_ULONG ulOperationType,
                                      CK_ULONG *pulFIPSStatus);

#include <dlfcn.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>

#include <stdlib.h>
#include <string.h>
#include <string_view>
#include <assert.h>
#include <vector>

#include <nss/nss.h>
#include <secmod.h>
#include <nspr.h>


#include <iostream>

void printInfo(const char* format, ...);
void printError(const char* format, ...);

char             g_error_message[200];
void*            g_pkilib = nullptr; 
CK_FUNCTION_LIST_PTR g_function_list = nullptr;

#define DLSYM_CK(name)   decltype(auto) name##_func = \
  (decltype(name) *)dlsym(g_pkilib, #name);


bool InitCK(const std::string& lib_path) {

  do {
    if(!g_pkilib) {
      g_pkilib = dlopen(lib_path.data(), RTLD_NOW);
    }

    if(!g_pkilib) {
      printError("failed to load driver lib");
      break;
    }

    CK_RV rv;
    // DLSYM_CK(C_GetFunctionList);
    // decltype(auto)
    decltype(C_GetFunctionList) *C_GetFunctionList_func =
        (decltype(C_GetFunctionList) *)dlsym(g_pkilib, "C_GetFunctionList");
    rv = C_GetFunctionList_func(&g_function_list);
    g_function_list->C_Initialize(NULL_PTR);

    // DLSYM_CK(C_Initialize);
    // rv = C_Initialize_func(NULL_PTR);
    if(CKR_OK != rv) {
      printError("Can not load PKCS#11 lib");
      break;
    }

    printInfo("load PKCS#11 lib succeed");
    return true;

  } while(0);

  UnInitCK();
  return false;
}


void UnInitCK()
{
  // if(g_function_list) {
  //   g_function_list->C_Finalize(NULL_PTR);
  //   g_function_list = nullptr;
  // }
  if(g_pkilib) {
    DLSYM_CK(C_Finalize)
    C_Finalize_func(NULL_PTR);
    dlclose(g_pkilib);
    g_pkilib = nullptr;
  }
}




CK_SESSION_HANDLE LoginCK(const std::string &pin)
{
  do {
    CK_SLOT_ID_PTR slot_ptr_list;
    CK_ULONG slot_count = 0;
    DLSYM_CK(C_GetSlotList)
    CK_RV rv = C_GetSlotList_func(CK_TRUE, NULL, &slot_count);
//    CK_RV rv = g_function_list->C_GetSlotList(CK_TRUE, NULL, &slot_count);
    if(rv != CKR_OK || slot_count <= 0) {
      printError("C_GetSlotList failed, please check your so file");
      break;
    }

    //return NULL_PTR;

    slot_ptr_list =  (CK_SLOT_ID_PTR)new CK_SLOT_ID[slot_count];
    rv = g_function_list->C_GetSlotList(CK_TRUE, slot_ptr_list, &slot_count);
    assert(rv == CKR_OK);

    CK_SESSION_HANDLE session = NULL_PTR;
    rv = g_function_list->C_OpenSession(slot_ptr_list[0],
                                   CKF_SERIAL_SESSION | CKF_RW_SESSION,
                                   NULL,
                                   NULL,
                                   &session);
    if(CKR_OK != rv || !session) {
      printError("C_OpenSession failed");
      break;
    }

    // rv = g_function_list->C_Login(session, CKU_USER, (CK_UTF8CHAR_PTR)pin.data(), pin.length());
    // if(CKR_OK != rv) {
    //   printMsg("C_Login failed");
    //   break;
    // }

    // CK_OBJECT_HANDLE private_key;
    // CK_OBJECT_CLASS ck_class = CKO_PRIVATE_KEY;
    // CK_ATTRIBUTE ck_template[] = {
    //     {CKA_CLASS, &ck_class, sizeof(ck_class)},
    //     {CKA_LABEL, (CK_VOID_PTR)"MyPrivateKey", strlen("MyPrivateKey")}
    // };

    // // 查找私钥
    // CK_ULONG object_count = 0;
    // if (g_function_list->C_FindObjectsInit(session, ck_template, 2) != CKR_OK ||
    //     g_function_list->C_FindObjects(session, &private_key, 1, &object_count) != CKR_OK ||
    //     g_function_list->C_FindObjectsFinal(session) != CKR_OK) {
    //   printMsg("Get Private Key failed");
    //   break;
    // }

    CK_OBJECT_HANDLE certificate;
    CK_OBJECT_CLASS ck_class1 = CKO_CERTIFICATE;
    bool is_token=true;
    CK_ATTRIBUTE ck_template1[] = {
        {CKA_CLASS, &ck_class1, sizeof(ck_class1)},
        {CKA_TOKEN, &is_token, sizeof(true)}
    };

    CK_ULONG object_count1 = 0;
    if (g_function_list->C_FindObjectsInit(session, ck_template1, 2) != CKR_OK ||
        g_function_list->C_FindObjects(session, &certificate, 1, &object_count1) != CKR_OK) {
      printError("Get certificate failed");
      break;
    }

    CK_ATTRIBUTE certificate_attr[] = {
        {CKA_VALUE, NULL, 0},  // 先查询长度
 //       {CKA_TOKEN, &IsToken, sizeof(IsToken)}
    };

    // g_function_list->C_FindObjectsInit(session, certificate_attr, 2);
    if(g_function_list->C_GetAttributeValue(session, certificate, certificate_attr, 1) != CKR_OK) {
      printError("Get size of certificate failed");
      break;
    }

    std::vector<CK_BYTE> cert_data(certificate_attr[0].ulValueLen);
    certificate_attr[0].pValue = cert_data.data();
    if(g_function_list->C_GetAttributeValue(session, certificate, certificate_attr, 1) != CKR_OK) {
      printError("Get bytes of certificate failed");
      break;
    }


    for(int j = 0; j < cert_data.size(); j++) {
        if(cert_data[j] == ' ')
            cert_data[j] = '_';
    }

    std::string cert_string(cert_data.size() * 2 + 1, 0);
    //char* cert_string = new char[cert_data.size() * 2 + 1]{0};
    for(int i = 0; i < cert_data.size(); i++) {
      //snprintf((char* )(cert_string.data() + i*2), 2, "%02x", (int)cert_data[i]);

        sprintf((char* )(cert_string.data() + i*2), "%02x", cert_data[i]);
    }

    g_function_list->C_FindObjectsFinal(session);


    int fd = open("cert.cer",O_RDWR|O_CREAT,0644);
    if(0 > fd) {
        printError("Open file fault");
        break ;
    }
    write(fd, cert_data.data(), cert_data.size());

    printInfo("Get bytes of certificate : %s\n", cert_string.data());
    return session;
  } while(0);

  UnInitCK();
  return NULL_PTR;
}

bool AddSoToNSSDatabase(const char* soPath) {
    // 1. 初始化 NSS 默认数据库（读写模式）
    if (NSS_InitReadWrite("sql:${HOME}/.pki/nssdb") != SECSuccess) {
        std::cerr << "NSS_InitReadWrite failed: " << PR_ErrorToString(PR_GetError(), PR_LANGUAGE_I_DEFAULT) << std::endl;
        return false;
    }

    // 2. 定义模块参数（替换为你的 .so 文件路径）
    SECMODModule* module = SECMOD_LoadUserModule(
        const_cast<char*>(soPath),  // .so 文件路径
        nullptr,                    // 模块名称（可选）
        0                           // 标志位（通常为 0）
    );

    if (!module || !module->loaded) {
        std::cerr << "Failed to load module: " << PR_ErrorToString(PR_GetError(), PR_LANGUAGE_I_DEFAULT) << std::endl;
        NSS_Shutdown();
        return false;
    }

    std::cout << "Module loaded successfully: " << module->commonName << std::endl;

    // 3. 释放模块资源（但不卸载）
    SECMOD_DestroyModule(module);

    // 4. 关闭 NSS
    NSS_Shutdown();
    return true;
}
